import os
import json
import sys
import argparse


def config():
	parser = argparse.ArgumentParser()
	parser.add_argument("--query_list", type=str, required=True)
	parser.add_argument("--stra_input", type=str, required=True)
	parser.add_argument("--base_input", type=str, default="outputs")
	parser.add_argument("--output_file", type=str, default="fake_doc_candidates")
	return parser.parse_args()


args = config()


# query list
with open(args.query_list) as f:
	query_list = []
	for line in f.read().strip().split("\n"):
		q, l = line.split("\t")
		if l == "-1":
			query_list.append(q)

print("Loss Query: {}".format("\t".join(query_list)))

# stra input
stra_input = {}
with open(args.stra_input) as f:
	lines = f.read().strip().split("\n")
	for line in lines:
		obj = json.loads(line)
		stra_input[obj["query"]["Q_query_basic"]] = obj


def load_base(filename):
	with open(filename) as f:
		content = f.read().strip().split("\n")

	table = []
	for line in content:
		table.append(line.split("\t"))

	header = table[0]
	query_type_col = header.index("query_type")
	url_col = header.index("url")
	title_col = header.index("title")

	video_query_type = "QuerySemanticDwellingTime"
	text_query_type = "DwellTimeTextVecRecallQueue"

	max_u_num = 50
	v_content = []
	t_content = []

	for row in table[1:]:
		if (row[query_type_col] == video_query_type) and (len(v_content) < max_u_num):
			v_content.append(row[title_col])
		elif (row[query_type_col] == text_query_type) and (len(t_content) < max_u_num):
			t_content.append(row[title_col])

	return v_content, t_content


def cst_fake_candidate(query, base_titles):
	stra_item = stra_input[query]
	first_candidate = stra_item["candidates_top"][0]["doc"]

	new_candidates = []
	for title in base_titles:
		new_candidate = first_candidate.copy()
		new_candidate["PU_title_basic"] = title
		new_candidates.append(new_candidate)
	return new_candidates


wf = open(args.output_file, "w", encoding="utf-8")
for qidx, query in enumerate(query_list):
	print("process {}.{}".format(qidx, query))

	try:
		b_v_cnt, b_t_cnt = load_base(os.path.join(args.base_input, "{}.txt".format(query)))
		base_cnt = b_v_cnt + b_t_cnt
		fake_candidates = cst_fake_candidate(query, base_cnt)
		for fc in fake_candidates:
			wf.write(json.dumps(fc, ensure_ascii=False))
			wf.write("\n")
	except Exception as e:
		print(e)
		continue
wf.close()